#include "hook.h"
#include <dlfcn.h>
#include "fiber.h"
#include "fd_manager.h"
#include "macro.h"
#include "timer.h"
#include "iomanager.h"
#include "log.h"
#include <stdarg.h>

static int g_tcp_connect_timeout = 5000;

// 是否可hook
static thread_local bool t_hook_enable = false;

#define HOOK_FUN(XX) \
    XX(sleep)        \
    XX(usleep)       \
    XX(nanosleep)    \
    XX(socket)       \
    XX(connect)      \
    XX(accept)       \
    XX(read)         \
    XX(readv)        \
    XX(recv)         \
    XX(recvfrom)     \
    XX(recvmsg)      \
    XX(write)        \
    XX(writev)       \
    XX(send)         \
    XX(sendto)       \
    XX(sendmsg)      \
    XX(close)        \
    XX(fcntl)        \
    XX(ioctl)        \
    XX(getsockopt)   \
    XX(setsockopt)

// hook_init() 放在一个静态对象的构造函数中调用，这表示在main函数运行之前
// 就会获取各个符号的地址并保存在全局变量中。
void hook_init()
{
    static bool is_inited = false;
    if (is_inited)
    {
        return;
    }
#define XX(name) name##_f = (name##_func)dlsym(RTLD_NEXT, #name);
    HOOK_FUN(XX);
#undef XX
}

// 上面的宏展开之后的效果如下：
// extern "C" {
//     sleep_fun sleep_f = nullptr; \
//     usleep_fun usleep_f = nullptr; \
//     ....
//     setsocketopt_fun setsocket_f = nullptr;
// };

// hook_init() {
//     ...
// 获取原API的地址
//     sleep_f = (sleep_fun)dlsym(RTLD_NEXT, "sleep"); \
//     usleep_f = (usleep_fun)dlsym(RTLD_NEXT, "usleep"); \
//     ...
//     setsocketopt_f = (setsocketopt_fun)dlsym(RTLD_NEXT, "setsocketopt");
// }

static uint64_t s_connect_timeout = -1;
struct _HookIniter
{
    _HookIniter()
    {
        // hook_init() 放在一个静态对象的构造函数中调用，这表示在main函数运行之前就会获取各个符号的地址并保存在全局变量中。
        hook_init();
        s_connect_timeout = g_tcp_connect_timeout;
    }
};

static _HookIniter s_hook_initer;

// 当前线程是否hook
bool is_hook_enable()
{
    return t_hook_enable;
}

// 设置当前线程的hook状态
void set_hook_enable(bool flag)
{
    t_hook_enable = flag;
}

struct timer_info
{
    int cancelled = 0;
};

template <typename OriginFun, typename... Args>
static ssize_t do_io(int fd, OriginFun fun, const char *hook_fun_name,
                     uint32_t event, int timeout_so, Args &&...args)
{
    if (!t_hook_enable)
    {
        return fun(fd, std::forward<Args>(args)...);
    }

    FdCtx::ptr ctx = FdMgr::GetInstance()->get(fd);
    if (!ctx)
    {
        return fun(fd, std::forward<Args>(args)...);
    }

    if (ctx->isClose())
    {
        errno = EBADF;
        return -1;
    }

    if (!ctx->isSocket() || ctx->getUserNonblock())
    {
        return fun(fd, std::forward<Args>(args)...);
    }

    uint64_t to = ctx->getTimeout(timeout_so);
    std::shared_ptr<timer_info> tinfo(new timer_info);

retry:
    ssize_t n = fun(fd, std::forward<Args>(args)...);
    while (n == -1 && errno == EINTR)
    {
        n = fun(fd, std::forward<Args>(args)...);
    }

    if (n == -1 && errno == EAGAIN)
    {
        IOManager *iom = IOManager::GetThis();
        Timer::ptr timer;
        std::weak_ptr<timer_info> winfo(tinfo);

        if (to != ((uint64_t)-1))
        {
            timer = iom->addConditionTimer(to, [winfo, fd, iom, event]()
                                           {
                auto t = winfo.lock();
                if(!t || t->cancelled) return;
                t->cancelled = ETIMEDOUT;
                iom->cancelEvent(fd,(IOManager::Event)(event)); }, winfo);
        }
        int ret = iom->addEvent(fd, (IOManager::Event)(event));
        if (SYLAR_UNLIKELY(ret))
        {
            LOG(ERROR, "%s addEvent(%d,%d)", hook_fun_name, fd, event);
            if (timer)
            {
                timer->cancel();
            }
            return -1;
        }
        else
        {
            Fiber::GetThis()->yield();
            if (timer)
            {
                timer->cancel();
            }
            if (tinfo->cancelled)
            {
                errno = tinfo->cancelled;
                return -1;
            }
            goto retry;
        }
    }
    return n;
}

extern "C"
{
#define XX(name) name##_func name##_f = nullptr;
    HOOK_FUN(XX);
#undef XX

    unsigned int sleep(unsigned int seconds)
    {
        if (!t_hook_enable)
        {
            return sleep_f(seconds);
        }

        Fiber::ptr fiber = Fiber::GetThis();
        IOManager *iom = IOManager::GetThis();
        // 将 IOManager::schedule 成员函数的地址转换为对应的成员函数指针类型，然后使用 std::bind 来绑定参数
        iom->addTimer(seconds * 1000, std::bind((void(Scheduler::*)(Fiber::ptr, int thread)) & IOManager::schedule, iom, fiber, -1));
        Fiber::GetThis()->yield();
        return 0;
    }

    int usleep(useconds_t usec)
    {
        if (!t_hook_enable)
        {
            usleep_f(usec);
        }
        Fiber::ptr fiber = Fiber::GetThis();
        IOManager *iom = IOManager::GetThis();
        iom->addTimer(usec / 1000, std::bind((void(Scheduler::*)(Fiber::ptr, int thread)) & IOManager::schedule, iom, fiber, -1));
        Fiber::GetThis()->yield();
        return 0;
    }

    int nanosleep(const struct timespec *req, struct timespec *rem)
    {
        if (!t_hook_enable)
        {
            nanosleep_f(req, rem);
        }

        int timeout_ms = req->tv_sec * 1000 + req->tv_nsec / 1000 / 1000;
        Fiber::ptr fiber = Fiber::GetThis();
        IOManager *iom = IOManager::GetThis();
        iom->addTimer(timeout_ms, std::bind((void(Scheduler::*)(Fiber::ptr, int thread)) & IOManager::schedule, iom, fiber, -1));
        Fiber::GetThis()->yield();
        return 0;
    }

    int socket(int domain, int type, int protocol)
    {
        if (!t_hook_enable)
        {
            socket_f(domain, type, protocol);
        }

        int fd = socket_f(domain, type, protocol);
        if (fd == -1)
            return fd;

        FdMgr::GetInstance()->get(fd, true);
        return fd;
    }

    int connect_with_timeout(int fd, const struct sockaddr *addr, socklen_t addrlen, uint64_t timeout_ms)
    {
        if (!t_hook_enable)
        {
            return connect_f(fd, addr, addrlen);
        }
        FdCtx::ptr ctx = FdMgr::GetInstance()->get(fd);
        if (!ctx || ctx->isClose())
        {
            errno = EBADF;
            return -1;
        }

        if (!ctx->isSocket())
        {
            return connect_f(fd, addr, addrlen);
        }
        if (ctx->getUserNonblock())
        {
            return connect_f(fd, addr, addrlen);
        }

        int n = connect_f(fd, addr, addrlen);
        if (n == 0)
        {
            return 0;
        }
        // EINPROGRESS操作正在进行
        else if (n != -1 || errno != EINPROGRESS)
        {
            return n;
        }

        IOManager *iom = IOManager::GetThis();
        Timer::ptr timer;
        std::shared_ptr<timer_info> tinfo(new timer_info);
        std::weak_ptr<timer_info> winfo(tinfo);

        if (timeout_ms != (uint64_t)-1)
        {
            timer = iom->addConditionTimer(timeout_ms, [winfo, fd, iom]
                                           {
                auto t = winfo.lock();
                if(!t || t->cancelled) return;

                // ETIMEDOUT连接超时
                t->cancelled = ETIMEDOUT;
                iom->cancelEvent(fd,IOManager::WRITE); }, winfo);
        }

        int ret = iom->addEvent(fd, IOManager::WRITE);
        if (ret == 0)
        {
            Fiber::GetThis()->yield();
            if (timer)
            {
                timer->cancel();
            }
            if (tinfo->cancelled)
            {
                errno = tinfo->cancelled;
                return -1;
            }
        }
        else
        {
            if (timer)
            {
                timer->cancel();
            }
            LOG(ERROR, "connect addEvent( %d,WRITE) error", fd);
        }

        int error = 0;
        socklen_t len = sizeof(int);
        if (-1 == getsockopt(fd, SOL_SOCKET, SO_ERROR, &error, &len))
        {
            return -1;
        }

        if (!error)
        {
            return 0;
        }
        else
        {
            errno = error;
            return -1;
        }
    }
}

int connect(int sockfd, const struct sockaddr *addr, socklen_t addrlen)
{
    return connect_with_timeout(sockfd, addr, addrlen, s_connect_timeout);
}

int accept(int sockfd, struct sockaddr *addr, socklen_t *addrlen)
{
    int fd = do_io(sockfd, accept_f, "accept", IOManager::READ, SO_RCVTIMEO, addr, addrlen);
    if (fd >= 0)
    {
        FdMgr::GetInstance()->get(fd, true);
    }
    return fd;
}

ssize_t read(int fd, void *buf, size_t count)
{
    return do_io(fd, read_f, "readv", IOManager::READ, SO_RCVTIMEO, buf, count);
}

ssize_t readv(int fd, const struct iovec *iov, int iovcnt)
{
    return do_io(fd, readv_f, "readv", IOManager::READ, SO_RCVTIMEO, iov, iovcnt);
}

ssize_t recv(int sockfd, void *buf, size_t len, int flags)
{
    return do_io(sockfd, recv_f, "recv", IOManager::READ, SO_RCVTIMEO, buf, len, flags);
}

ssize_t recvfrom(int sockfd, void *buf, size_t len, int flags, struct sockaddr *src_addr, socklen_t *addrlen)
{
    return do_io(sockfd, recvfrom_f, "recvfrom", IOManager::READ, SO_RCVTIMEO, buf, len, flags, src_addr, addrlen);
}

ssize_t recvmsg(int sockfd, struct msghdr *msg, int flags)
{
    return do_io(sockfd, recvmsg_f, "recvmsg", IOManager::WRITE, SO_SNDTIMEO, msg, flags);
}

ssize_t write(int fd, const void *buf, size_t count)
{
    return do_io(fd, write_f, "write", IOManager::WRITE, SO_SNDTIMEO, buf, count);
}

ssize_t writev(int fd, const struct iovec *iov, int iovcnt)
{
    return do_io(fd, write_f, "writev", IOManager::WRITE, SO_SNDTIMEO, iov, iovcnt);
}

ssize_t send(int fd, const void *msg, size_t len, int flags)
{
    return do_io(fd, send_f, "send", IOManager::WRITE, SO_SNDTIMEO, msg, len, flags);
}

ssize_t sendto(int fd, const void *msg, size_t len, int flags, const struct sockaddr *to, socklen_t tolen)
{
    return do_io(fd, sendto_f, "sendto", IOManager::WRITE, SO_SNDTIMEO, msg, len, flags, to, tolen);
}

ssize_t sendmsg(int fd, const struct msghdr *msg, int flags)
{
    return do_io(fd, sendmsg_f, "sendmsg", IOManager::WRITE, SO_SNDTIMEO, msg, flags);
}

int close(int fd)
{
    if (!t_hook_enable)
    {
        return close(fd);
    }

    FdCtx::ptr ctx = FdMgr::GetInstance()->get(fd);
    if (ctx)
    {
        auto iom = IOManager::GetThis();
        if (iom)
        {
            iom->cancelAll(fd);
        }
        FdMgr::GetInstance()->del(fd);
    }
    return close(fd);
}

int fnctl(int fd, int cmd, ... /* arg*/)
{
    va_list va;
    va_start(va, cmd);

    switch (cmd)
    {
    case F_SETFL:
    {
        int arg = va_arg(va, int);
        va_end(va);

        FdCtx::ptr ctx = FdMgr::GetInstance()->get(fd);
        if (!ctx || ctx->isClose() || !ctx->isSocket())
        {
            return fcntl(fd, cmd, arg);
        }
        ctx->setUserNonblock(arg & O_NONBLOCK);
        if (ctx->getSysNonblock())
        {
            arg |= O_NONBLOCK;
        }
        else
        {
            arg &= ~O_NONBLOCK;
        }
    }
    break;
    case F_GETFL:
    {
        va_end(va);
        int arg = fcntl_f(fd, cmd);
        FdCtx::ptr ctx = FdMgr::GetInstance()->get(fd);
        if (!ctx || ctx->isClose() || !ctx->isSocket())
        {
            return arg;
        }
        if (ctx->getUserNonblock())
        {
            return arg | O_NONBLOCK;
        }
        else
        {
            return arg & ~O_NONBLOCK;
        }
    }
    break;
    case F_DUPFD:
    case F_DUPFD_CLOEXEC:
    case F_SETFD:
    case F_SETOWN:
    case F_SETSIG:
    case F_SETLEASE:
    case F_NOTIFY:
#ifdef F_SETPIPE_SZ
    case F_SETPIPE_SZ:
#endif
    {
        int arg = va_arg(va, int);
        va_end(va);
        return fcntl_f(fd, cmd, arg);
    }
    break;
    case F_GETFD:
    case F_GETOWN:
    case F_GETSIG:
    case F_GETLEASE:
#ifdef F_GETPIPE_SZ
    case F_GETPIPE_SZ:
#endif
    {
        va_end(va);
        return fcntl_f(fd, cmd);
    }
    break;
    case F_SETLK:
    case F_SETLKW:
    case F_GETLK:
    {
        struct flock *arg = va_arg(va, struct flock *);
        va_end(va);
        return fcntl_f(fd, cmd, arg);
    }
    break;
    case F_GETOWN_EX:
    case F_SETOWN_EX:
    {
        struct f_owner_exlock *arg = va_arg(va, struct f_owner_exlock *);
        va_end(va);
        return fcntl_f(fd, cmd, arg);
    }
    break;
    default:
        va_end(va);
        return fcntl_f(fd, cmd);
    }
}

int ioctl(int fd, unsigned long int request, ...)
{
    va_list va;
    va_start(va, request);

    void *arg = va_arg(va, void *);
    va_end(va);

    if (FIONBIO == request)
    {
        bool user_nonblock = !!*(int *)arg;
        FdCtx::ptr ctx = FdMgr::GetInstance()->get(fd);
        if (!ctx || ctx->isClose() || !ctx->isSocket())
        {
            return ioctl_f(fd, request, arg);
        }
        ctx->setUserNonblock(user_nonblock);
    }

    return ioctl_f(fd, request, arg);
}

int getsockopt(int sockfd, int level, int opname, void *optval, socklen_t *optlen)
{
    return getsockopt_f(sockfd, level, opname, optval, optlen);
}

int setsockopt(int sockfd, int level, int optname, const void *optval, socklen_t optlen)
{
    if (!t_hook_enable)
    {
        return setsockopt(sockfd, level, optname, optval, optlen);
    }

    if (level == SOL_SOCKET)
    {
        if (optname == SO_RCVTIMEO || optname == SO_SNDTIMEO)
        {
            FdCtx::ptr ctx = FdMgr::GetInstance()->get(sockfd);
            if (ctx)
            {
                const timeval *v = (const timeval *)optval;
                ctx->setTimeout(optname, v->tv_sec * 1000 + v->tv_usec / 1000);
            }
        }
    }

    return setsockopt_f(sockfd, level, optname, optval, optlen);
}
